/* GEMM 优化示例
核心功能：
 实现高性能矩阵乘法 C = A * B，其中：
 - A矩阵维度：M x K
 - B矩阵维度：K x N
 - C矩阵维度：M x N

CUDA编程流程：
 1. 主机端内存分配与初始化
 2. 设备端显存分配 (cudaMalloc)
 3. 主机到设备数据传输 (cudaMemcpy)
 4. 核函数执行配置（网格/线程块维度）
 5. 启动核函数 (cudaLaunchKernelGGL)
 6. 设备到主机结果回传 (cudaMemcpy)
 7. 资源释放 (cudaFree/cudaEventDestroy)
*/

#include <iostream>
#include <math.h>
#include <cuda_runtime.h>

#define DEFAULT_M 4096
#define DEFAULT_N 4096
#define DEFAULT_K 4096
#define DEFAULT_BLOCK_X 8
#define DEFAULT_BLOCK_Y 8
#define CHECK_OUTPUT_CORRECT false

#define CHECK_HIP(cmd) { cudaError_t error = cmd; if (error != cudaSuccess) { \
    printf("HIP error: %s:%d '%s'\n", __FILE__,__LINE__,cudaGetErrorString(error)); \
    exit(EXIT_FAILURE);} }

// 矩阵乘法核函数
__launch_bounds__(1024) __global__ void matrixMul(const float* __restrict__ A,
                         const float* __restrict__ B,
                         float* __restrict__ C,
			 long M,
			 long N,
			 long K,
                         int block_x,
			 int block_y) 
{
    const int row = blockIdx.y * block_y + threadIdx.y;
    const int col = blockIdx.x * block_x + threadIdx.x;
    if (row < M && col < N) {
        float sum = 0.0f;
        for (int k = 0; k < K; ++k) {
            sum += A[row * K + k] * B[k * N + col];
        }
        C[row * N + col] = sum;
    }
}

void cpuMatrixMul(const float* A, const float* B, float* C, long M, long N, long K) {
    #pragma omp parallel for collapse(2)
    for (int row = 0; row < M; ++row) {
        for (int col = 0; col < N; ++col) {
            float sum = 0.0f;
            for (int k = 0; k < K; ++k) {
                sum += A[row*K + k] * B[k*N + col];
            }
            C[row*N + col] = sum;
        }
    }
}

int main(int argc, char** argv) {
    // 参数解析
    int block_x = DEFAULT_BLOCK_X;
    int block_y = DEFAULT_BLOCK_Y;
    long M = DEFAULT_M;
    long N = DEFAULT_N;
    long K = DEFAULT_K;
    bool use_default = true;
    if(argc == 6) {
	M = atoi(argv[1]);
	N = atoi(argv[2]);
	K = atoi(argv[3]);
        block_x = atoi(argv[4]);
        block_y = atoi(argv[5]);
        use_default = false;
    } else {
	std::cout << "程序格式：./exe M N K block_x block_y\n";
        std::cout << "警告：需要五个分块参数（M, N, K, block_x ,block_y），已使用默认值\n";
    }   
    
    // 参数校验
    if(M * N * K != 4096LL * 4096LL * 4096LL || 
       block_x * block_y > 1024        || 
       block_x > 1024 || block_x < 1   ||
       block_y > 1024 || block_y < 1    ) 
    {
        std::cout << "错误：M N K 不符合要求或 block_x block_y 分块尺寸超过设备限制，重置为默认值\n";
	std::cout << "需满足 M * N * K < 4096 * 4096 * 4096, block_x * block_y < 1024\n";
	M = DEFAULT_M;
	N = DEFAULT_N;
	K = DEFAULT_K;
        block_x = DEFAULT_BLOCK_X;
        block_y = DEFAULT_BLOCK_Y;
        use_default = true;
    }

    // 打印配置信息
    if(use_default) {
        std::cout << "提示：使用默认M N K 和分块大小\n";
    }
    
    std::cout << "\nGEMM配置参数：\n"
              << "矩阵维度：\tM=" << M << "\tN=" << N << "\tK=" << K << "\n"
              << "分块尺寸：\tBLOCK_X=" << block_x << "\tBLOCK_Y=" << block_y << "\n"
              << "网格维度：\tGRID_X=" << (N+block_x-1)/block_x << "\tGRID_Y=" 
              << (M+block_y-1)/block_y << "\n\n";

    // 主机内存分配
    float *h_A = new float[M*K];
    float *h_B = new float[K*N];
    float *h_C = new float[M*N];
    float* h_C_ref = new float[M*N];

    // 矩阵初始化
    #pragma omp parallel for
    for (int i = 0; i < M*K; ++i) h_A[i] = (rand() % 255) / 255.0;
    #pragma omp parallel for
    for (int i = 0; i < K*N; ++i) h_B[i] = (rand() % 255) / 255.0;;

    // 设备显存分配
    float *d_A, *d_B, *d_C;
    CHECK_HIP(cudaMalloc(&d_A, M*K*sizeof(float)));
    CHECK_HIP(cudaMalloc(&d_B, K*N*sizeof(float)));
    CHECK_HIP(cudaMalloc(&d_C, M*N*sizeof(float)));

    // 数据传输：主机→设备
    CHECK_HIP(cudaMemcpy(d_A, h_A, M*K*sizeof(float), cudaMemcpyHostToDevice));
    CHECK_HIP(cudaMemcpy(d_B, h_B, K*N*sizeof(float), cudaMemcpyHostToDevice));

    // 核函数配置
    dim3 block(block_x, block_y);
    dim3 grid((N + block_x-1)/block_x, (M + block_y-1)/block_y);
    size_t shared_size = (block_y * block_x + block_x * block_x) * sizeof(float);

    cudaEvent_t start, stop;
    CHECK_HIP(cudaEventCreate(&start));
    CHECK_HIP(cudaEventCreate(&stop));

    // 执行核函数
    CHECK_HIP(cudaEventRecord(start));
    matrixMul<<<grid, block, shared_size>>>(d_A, d_B, d_C, M, N, K,block_x, block_y);
    CHECK_HIP(cudaEventRecord(stop));
    CHECK_HIP(cudaEventSynchronize(stop));

    float milliseconds = 0;
    CHECK_HIP(cudaEventElapsedTime(&milliseconds, start, stop));
    std::cout << "Kernel time: " << milliseconds << " ms\n";

    // 结果回传
    CHECK_HIP(cudaMemcpy(h_C, d_C, M*N*sizeof(float), cudaMemcpyDeviceToHost));
    
    if(CHECK_OUTPUT_CORRECT) {
    	cpuMatrixMul(h_A, h_B, h_C_ref, M, N, K);
    	float max_error = 0.0f;
    	#pragma omp parallel for reduction(max:max_error)
    	for (int i = 0; i < M*N; i=i+100) {
        	float ref = h_C_ref[i];
        	float val = h_C[i];
        	if (ref != 0) {
            		float error = fabs(val - ref)/fabs(ref);
            		max_error = fmax(max_error, error);
        	} else if (fabs(val) > 1e-5) {
            	max_error = 1.0f; 
            }
    	}    
    const float tolerance = 1e-3f;
    bool correct = (max_error <= tolerance);
    std::cout << "\n最大相对误差: " << max_error 
              << " 允许误差: " << tolerance
              << "\n验证结果: " << (correct ? "通过" : "失败") << "\n";    
    }
    delete[] h_A;
    delete[] h_B;
    delete[] h_C;
    delete[] h_C_ref;
    CHECK_HIP(cudaFree(d_A));
    CHECK_HIP(cudaFree(d_B));
    CHECK_HIP(cudaFree(d_C));
    CHECK_HIP(cudaEventDestroy(start));
    CHECK_HIP(cudaEventDestroy(stop));
    return 0;
}
